#include <iostream>
#include <fstream>
#include <cmath>
#include <vector>
#include <cstdlib> 

using namespace std;

// 参数设置
const int n = 64; // 网格点数
const double pi = 3.141592653589793;

// 计算lambda_k
double lambda_k(int k, double w) {
    return 1 - 2 * w * pow(sin((k * pi) / (2 * 64)), 2);
}

// 计算所需的迭代次数
int get_n(int k, double w) {
    double lk = lambda_k(k, w);
    if (abs(lk) < 1e-10) return 100; // 避免除以零
    return static_cast<int>(min(100.0, (log(0.01) / log(abs(lk)))));
}

int main() {
    vector<int> X;
    vector<int> Y1, Y2;

    for (int i = 1; i < n; ++i) {
        X.push_back(i);
        Y1.push_back(get_n(i, 1.0));
        Y2.push_back(get_n(i, 2.0 / 3.0));
    }

    // 将数据写入文件
    ofstream dataFile("iterations_data.dat");
    if (dataFile.is_open()) {
        dataFile << "# k omega=1 omega=2/3" << endl;
        for (size_t i = 0; i < X.size(); ++i) {
            if(i==31){
                Y1[i]=0;}
            dataFile << X[i] << " " << Y1[i] << " " << Y2[i] << endl;
        }
        dataFile.close();
    }

    // 使用gnuplot绘制图形
    const char* gnuplot_script = R"(
        set terminal pngcairo size 600,600 enhanced font 'Verdana,10'
        set output 'iterations_a.png'
        set title 'Iterations for {\omega}=1'
        set xlabel 'Wavenumber k'
        set ylabel 'Iterations'
        set grid
        unset key
        plot 'iterations_data.dat' using 1:2 with lines lt 1 lw 2 lc 'black'
        set output 'iterations_b.png'
        set title 'Iterations for {\omega}=2/3'
        plot 'iterations_data.dat' using 1:3 with lines lt 1 lw 2 lc 'black'
    )";

    // 将gnuplot脚本写入临时文件
    ofstream gnuplotFile("temp.gnuplot");
    if (gnuplotFile.is_open()) {
        gnuplotFile << gnuplot_script;
        gnuplotFile.close();

        // 调用gnuplot执行脚本
        system("gnuplot temp.gnuplot");

        // 删除临时文件
        remove("temp.gnuplot");
    }

    cout << "图片生成完成，文件名为 iterations_a.png 和 iterations_b.png。" << endl;

    return 0;
}